# python-stix
from stix.core import STIXPackage
from stix.common.vocabs import IndicatorType
from stix.indicator import Indicator
import logging
import traceback
import argparse
import os
from librule import Type, LibruleWrapper
from common_dl import *
import sys
import json
import stix.utils.parser


class StixSo(object):

    def __init__(self):

        self.type_supported = [Type.IP,Type.URL,Type.SHA1,Type.DOMAIN]

        self.so_valid = self.create_so_object()
        self.so_invalid = self.create_so_object()
        self.so_duplicate = self.create_so_object()

        self.OP_APPLY_CONDITION='ANY'
        self.OP_CONDITION='Equals'
        self.OP_OPERATOR='OR'

        self.librule = LibruleWrapper()
        self._db_con = create_db_connection()

        self.stix_xml = None

    def create_so_object(self):
        so={}
        for i in self.type_supported:
            so[i] = []
        return so

    def push_so(self,so,data,type):

        if type not in self.type_supported:
            logging.error('The type {} is not unsupported'.format(type))
            raise codeException(Code.UNSUPPORT_TYPE)

        if isinstance(data, list):
            so[type].extend(data)
        else:
            so[type].append(data)

    def _check_condition(self,value):
        condition, apply = False, False

        logging.debug('checking condition: {}'.format(value.condition))
        if value.condition == self.OP_CONDITION:
            condition = True
        if value.apply_condition:
            logging.debug('checking apply_condition: {}'.format(value.apply_condition))
            if value.apply_condition == self.OP_APPLY_CONDITION:
                apply = True
        else:
            apply = True

        return condition and apply

    def _has_attribute(self,observable,attr):
        return hasattr(observable,attr) and getattr(observable,attr) is not None

    def __get_observable(self,observable):

        if self._has_attribute(observable,'object_'):
            logging.debug('Find Single model')
            return [observable]
        elif self._has_attribute(observable,'observable_composition'):
            logging.debug('Find Composition model')
            if observable.observable_composition.operator == self.OP_OPERATOR:
                return observable.observable_composition.observables
        logging.warn('Unsupported model')
        raise codeException(Code.INVALID_STIX_FORMAT)

    def _parse_composite_indicator_expression(self,indicator):
        #only support in stix Version: 1.2.0.1
        ob =[]
        if indicator.composite_indicator_expression.operator != self.OP_OPERATOR:
            logging.warn('Unsupported composite_indicator_expression operator')
            raise codeException(Code.INVALID_STIX_FORMAT)

        if not self._has_attribute(indicator.composite_indicator_expression,'indicator'):
            ind = indicator.composite_indicator_expression.to_dict()['indicators']

        else:
            ind = indicator.composite_indicator_expression.indicator

        for i in ind:
            ob.append(Indicator.from_dict(i).observable)
        return ob

    def load_stix(self):

        logging.debug('Try to load STIX {} xml'.format(self.stix_xml))
        while True:
            try:
                import warnings
                warnings.filterwarnings("ignore")
                package = STIXPackage.from_xml(self.stix_xml)
            except stix.utils.parser.UnsupportedVersionError as e:
                logging.warn('Ignore version and try again.{}'.format(e))
                parser = stix.utils.parser.EntityParser()
                return parser.parse_xml(self.stix_xml, check_version=False)
            except:
                logging.error(traceback.format_exc())
                raise codeException(Code.INVALID_STIX_FORMAT)
            else:
                return package

    def parse_stix(self):

        package = self.load_stix()
        #
        #print all stix
        # print package.to_dict()

        so = self.create_so_object()
        for indicator in package.indicators:
            ind_type = str(indicator.indicator_types[0])
            logging.debug('Watchlist Type: {}'.format(ind_type))

            #
            # print this indicator
            # print indicator.to_dict()
            #

            if self._has_attribute(indicator,'composite_indicator_expression'):
                observables= self._parse_composite_indicator_expression(indicator)
            else:
                observables = self.__get_observable(indicator.observable)

            #
            #  loop  observables to get data
            #

            for observable in observables:
                if ind_type == IndicatorType.TERM_URL_WATCHLIST:
                    v=observable.object_.properties.value
                    if self._check_condition(v):
                        self.push_so(so,v.value,Type.URL)
                        logging.debug('Find {}: {}'.format(ind_type,v.value))
                elif ind_type == IndicatorType.TERM_DOMAIN_WATCHLIST:
                    v=observable.object_.properties.value
                    if self._check_condition(v):
                        self.push_so(so,v.value,Type.DOMAIN)
                        logging.debug('Find {}: {}'.format(ind_type,v.value))
                elif ind_type == IndicatorType.TERM_IP_WATCHLIST:
                    v=observable.object_.properties.address_value
                    if self._check_condition(v):
                        self.push_so(so,v.value,Type.IP)
                        logging.debug('Find {}: {}'.format(ind_type,v.value))
                elif ind_type == IndicatorType.TERM_FILE_HASH_WATCHLIST:
                    for h in observable.object_.properties.hashes:
                        v = h.simple_hash_value
                        if self._check_condition(v):
                            self.push_so(so,v.value,Type.SHA1)
                            logging.debug('Find {}: {}'.format(ind_type,v.value))
        return so

    def __load_data_from_db(self):
        self._db_con.execute('SELECT data,type from tbl_deny_list ORDER BY type;')
        result = self._db_con.fetchall()

        so = self.create_so_object()
        for data, type in result:
            so[type].append(data)
        self.librule.import_rules(so)

    def _is_duplicate_self(self, data, rule_type):
        # # use another librule handle
        # itself = LibruleWrapper()
        # itself.import_rules(self.so_valid)
        # return itself.match_all(data,rule_type)
        return data in self.so_valid[rule_type]

    def _is_duplicate_db(self, data, rule_type):
        return self.librule.match_all(data,rule_type)

    def is_duplicate(self, data, rule_type):
        return self._is_duplicate_self(data,rule_type) or self._is_duplicate_db(data,rule_type)

    def do_valid_duplicate(self, so):

        cnt = 0
        self.__load_data_from_db()

        logging.debug('Do valid/deDuplication check ...')
        for rule_type,rule_data in so.iteritems():
            for data in rule_data:
                data = data.strip()
                if self.librule.is_valid_all(data,rule_type) and is_valid_length(data):
                    if self.is_duplicate(data,rule_type):
                        logging.warn('Find duplicate so {}'.format(data))
                        self.push_so(self.so_duplicate,data,rule_type)
                    else:
                        if cnt > MAX_DENYLIST_CNT:
                            logging.error('So count > {}'.format(MAX_DENYLIST_CNT))
                            codeException(Code.TOO_MANY_ITEM)

                        logging.debug('Find normal so {}'.format(data))
                        self.push_so(self.so_valid,data,rule_type)
                        cnt += 1
                else:
                    logging.warn('Find invalid so {}'.format(data))
                    self.push_so(self.so_invalid,data,rule_type)

    def _format_result(self,code):

        ret = {}
        ret['code'] = code
        if code == Code.DATA_VALIDATE_DE_DUPLICATE:

            ret['value'] = {}
            ret['value']['valid']=[]
            ret['value']['invalid']=[]
            for rule_type,rule_data in self.so_valid.iteritems():
                for i in rule_data:
                    ret['value']['valid'].append({'data':i,'type':rule_type})
            for rule_type,rule_data in self.so_invalid.iteritems():
                for i in rule_data:
                    ret['value']['invalid'].append({'reason':Code.REASON_INVALID,'data':i,'type':rule_type})
            for rule_type,rule_data in self.so_duplicate.iteritems():
                for i in rule_data:
                    ret['value']['invalid'].append({'reason':Code.REASON_DUPLICATE,'data':i,'type':rule_type})

        result = json.dumps(ret)
        logging.debug('Result:{}'.format(result))
        return result

    def cmd_parse_denylist(self,stix_xml):

        self.stix_xml = stix_xml
        try:
            original_so = self.parse_stix()
            self.do_valid_duplicate(original_so)
        except codeException:
            return self._format_result(Code.INVALID_STIX_FORMAT)
        except:
            logging.error(traceback.format_exc())
            return self._format_result(Code.OTHER_ERROR)
        return self._format_result(Code.DATA_VALIDATE_DE_DUPLICATE)

def set_arg():
    parser = argparse.ArgumentParser('stix')
    parser.add_argument('--parse',help='Get SO from STIX format Watchlist')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    return parser.parse_args()

if __name__ == '__main__':
    arg = set_arg()
    log_init()
    if not os.path.exists(arg.parse):
        logging.error('No such file {}'.format(arg.parse))
        sys.exit(1)

    ss = StixSo()
    print ss.cmd_parse_denylist(arg.parse)
    log_destr()

